import os
import torch
import torchvision
import torchvision.transforms as transforms
from torch import nn
from torch.utils.data import DataLoader
import numpy as np
import matplotlib.pyplot as plt
from torchvision.utils import save_image


# ========================= 数据加载与预处理 =========================
def load_data(batch_size=64):
    transform = transforms.ToTensor()
    train_data = torchvision.datasets.MNIST(root='data', train=True, download=True, transform=transform)
    test_data = torchvision.datasets.MNIST(root='data', train=False, download=True, transform=transform)
    train_dataloader = DataLoader(dataset=train_data, batch_size=batch_size)
    test_dataloader = DataLoader(dataset=test_data, batch_size=batch_size)
    return train_dataloader, test_dataloader


# ========================= 模型定义 =========================
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(1, 6, 3, stride=1, padding=1),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(6, 16, 5, stride=1, padding=1),
            nn.MaxPool2d(2, 2)
        )
        self.fc = nn.Sequential(
            nn.Linear(576, 120),
            nn.Linear(120, 84),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


# ========================= 训练函数 =========================
def train_model(network, train_dataloader, device, loss_fn, optimizer, epochs=10):
    network.train()
    losses = []
    for epoch in range(epochs):
        loss_sum = 0
        for X, y in train_dataloader:
            X, y = X.to(device), y.to(device)
            pred = network(X)
            loss = loss_fn(pred, y)
            loss_sum += loss.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        mean_loss = loss_sum / len(train_dataloader.dataset)
        losses.append(mean_loss)
        print(f"Epoch {epoch + 1} loss: {mean_loss:.6f}")
    torch.save(network.state_dict(), "model.pth")
    plt.xlabel("Epochs")
    plt.ylabel("Loss Value")
    plt.plot(range(epochs), losses)
    plt.show()


# ========================= 测试函数 =========================
def evaluate_model(network, test_dataloader, device):
    network.eval()
    positive, negative = 0, 0
    with torch.no_grad():
        for X, y in test_dataloader:
            X, y = X.to(device), y.to(device)
            pred = network(X)
            for item in zip(pred, y):
                if torch.argmax(item[0]) == item[1]:
                    positive += 1
                else:
                    negative += 1
    acc = positive / (positive + negative)
    print(f"Accuracy: {acc * 100:.2f}%")
    return acc


# ========================= FGSM攻击与可视化 =========================
def fgsm_attack(network, test_dataloader, device, loss_fn, epsilons, output_dir="adversarial_samples"):
    network.eval()
    acc_list = []

    # 创建输出文件夹
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    for epsilon in epsilons:
        positive, negative = 0, 0
        for batch_idx, (X, y) in enumerate(test_dataloader):
            X, y = X.to(device), y.to(device)
            X.requires_grad = True
            pred = network(X)
            loss = loss_fn(pred, y)
            network.zero_grad()
            loss.backward()
            X_adv = torch.clamp(X + epsilon * X.grad.sign(), 0, 1)
            pred = network(X_adv)
            for item in zip(pred, y):
                if torch.argmax(item[0]) == item[1]:
                    positive += 1
                else:
                    negative += 1

             # 保存原始图像和对抗样本
            if batch_idx == 0 and epsilon == epsilons[0]:  # 仅在第一个 epsilon 时保存原始图像
                original_images_path = os.path.join(output_dir, "original_images.png")
                save_image(X.cpu(), original_images_path)

            if batch_idx == 0:  # 每个 epsilon 保存一次对抗样本
                adversarial_images_path = os.path.join(output_dir, f"adversarial_epsilon_{epsilon:.2f}.png")
                save_image(X_adv.cpu(), adversarial_images_path)
        acc = positive / (positive + negative)
        acc_list.append(acc)
        print(f"epsilon={epsilon} acc: {acc * 100:.2f}%")

    # 绘制并保存 epsilon 与准确率的关系图
    acc_plot_path = os.path.join(output_dir, "accuracy_vs_epsilon.png")
    plt.xlabel("Epsilon")
    plt.ylabel("Accuracy")
    plt.ylim(0, 1)
    plt.plot(epsilons, acc_list, marker='o')
    plt.savefig(acc_plot_path)
    print(f"Accuracy plot saved to {acc_plot_path}")


# ========================= 主函数 =========================
def main():
    # 加载数据
    batch_size = 64
    train_dataloader, test_dataloader = load_data(batch_size)

    # 设置设备
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")

    # 初始化模型、损失函数和优化器
    network = LeNet().to(device)
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(network.parameters(), lr=0.001, momentum=0.9)

    # 训练或加载模型
    if os.path.exists('model.pth'):
        network.load_state_dict(torch.load('model.pth'))
        print("Loaded pre-trained model.")
    else:
        print("Training model...")
        train_model(network, train_dataloader, device, loss_fn, optimizer)

    # 测试模型
    print("Evaluating model...")
    evaluate_model(network, test_dataloader, device)

    # FGSM攻击与可视化
    epsilons = [0.01, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45]
    print("Performing FGSM attack...")
    fgsm_attack(network, test_dataloader, device, loss_fn, epsilons)


if __name__ == "__main__":
    main()
